import cv2
import numpy as np
import matplotlib.pyplot as plt


def gray_translation(image_path, save_plot_path, save_path=None, show_result=True):
    # 读取图像
    img = cv2.imread(image_path, 0)
    if img is None:
        raise ValueError("无法读取图像，请检查路径")

    # 计算整张图的平均值
    global_mean = np.mean(img)

    height, width = img.shape
    num_regions = height
    region_height = 1
    corrected_img = np.copy(img)

    for i in range(num_regions):
        y1 = i * region_height
        y2 = (i + 1) * region_height if i != num_regions - 1 else height
        region = img[y1:y2, :]
        region_mean = np.mean(region)
        corrected_img[y1:y2, :] = img[y1:y2, :] - region_mean + global_mean

    corrected_img = np.clip(corrected_img, 0, 255).astype(np.uint8)

    # 结果展示与保存
    if show_result:
        plt.rcParams['font.sans-serif'] = ['Microsoft YaHei'] 
        plt.figure(figsize=(10, 5))
        plt.subplot(121), plt.imshow(img, cmap='gray'), plt.title('原始图像')
        plt.subplot(122), plt.imshow(corrected_img, cmap='gray'), plt.title('灰度平移后图像')
        plt.savefig(save_plot_path)
        plt.show()
    
    if save_path:
        cv2.imwrite(save_path, corrected_img)
    
    return corrected_img


if __name__ == "__main__":
    # 输入图像路径（请替换为你的图像路径）
    image_path = "./gray_translation/origin.png"
    # 输出路径（可选）
    save_path = "./gray_translation/corrected_image.png"
    compare_path = "./gray_translation/compare_image.png"

    # 执行灰度矫正
    corrected_img = gray_translation(image_path, save_path, compare_path)